package com.timeriver.feature_project

import org.apache.spark.ml.feature.{OneHotEncoderEstimator, OneHotEncoderModel, StringIndexer}
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
  * 独热向量编码：将字符串变量转为二进制编码
  *   1. 字符串先转为数值类型
  *   2. 数值类型再转为二进制类型
  */
object OneHotEncoderDemo {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .appName("皮尔逊独立检验计算")
      .master("local[6]")
      .getOrCreate()

    val df = session.createDataFrame(
      Seq((0, "a"), (1, "b"), (2, "c"), (3, "a"), (4, "a"), (5, "c"))
    ).toDF("id", "category")

    /** 字符串转数值 */
    val indexer = new StringIndexer()
      .setInputCol("category")
      .setOutputCol("categoryIndex")

    val frame: DataFrame = indexer.fit(df).transform(df)
    frame.show(false)

    /** 独热向量编码 */
    val estimator: OneHotEncoderEstimator = new OneHotEncoderEstimator()
      .setInputCols(Array("categoryIndex"))
      .setOutputCols(Array("categoryVector"))

    val model: OneHotEncoderModel = estimator.fit(frame)

    val res: DataFrame = model.transform(frame)
    res.show(false)
  }
}
